package com.example.studyllm.component.ai.model.service;

import cn.hutool.core.util.IdUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.extra.spring.SpringUtil;
import com.example.studyllm.bo.ChatBO;
import com.example.studyllm.bo.ModelBO;
import com.example.studyllm.component.ai.ChatAgent;
import com.example.studyllm.component.ai.model.ChatModelBuilder;
import com.example.studyllm.component.constants.CommonConst;
import com.example.studyllm.component.utils.PromptUtil;
import com.example.studyllm.config.exception.BusinessException;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.service.TokenStream;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import dev.langchain4j.store.memory.chat.InMemoryChatMemoryStore;
import org.springframework.stereotype.Service;

@Service
public class ModelFactory {
    private ChatMemoryStore memoryStore =  new InMemoryChatMemoryStore();

    public ChatLanguageModel buildChatModel(ModelBO modelBO){
        ChatLanguageModel model = SpringUtil.getBean(modelBO.getCode());
        if(model==null){
            throw new BusinessException("模型不存在");
        }
        return model ;
    }

    public StreamingChatLanguageModel buildChatModelStream(ModelBO modelBO){
        StreamingChatLanguageModel streamingChatLanguageModel =  SpringUtil.getBean(modelBO.getCode()+ CommonConst.BEAN_PREFIX_STREAMCHAT);
        if(streamingChatLanguageModel==null){
            throw new BusinessException("模型不存在");
        }
        return streamingChatLanguageModel ;
    }

    private AiServices<ChatAgent> build(StreamingChatLanguageModel streamModel, ChatLanguageModel model, ChatBO req) {
        AiServices<ChatAgent> aiServices = AiServices.builder(ChatAgent.class)
                .chatMemoryProvider(memoryId -> MessageWindowChatMemory.builder()
                        .id(req.getConversationId())
                        .chatMemoryStore(memoryStore)
                        .maxMessages(req.getMaxMessageSize())
                        .build());
        if (StrUtil.isNotBlank(req.getPromptText())) {
            aiServices.systemMessageProvider(memoryId -> req.getPromptText());
        }
        if (streamModel != null) {
            aiServices.streamingChatLanguageModel(streamModel);
        }
        if (model != null) {
            aiServices.chatLanguageModel(model);
        }
        return aiServices;
    }

    public TokenStream chatStream(ModelBO modelBO, ChatBO chatBO){
        StreamingChatLanguageModel model =buildChatModelStream(modelBO);
        if (StrUtil.isBlank(chatBO.getConversationId())) {
            chatBO.setConversationId(IdUtil.simpleUUID());
        }

        ChatAgent agent = build(model, null, chatBO).build();
        if (chatBO.getPrompt() == null) {
            chatBO.setPrompt(PromptUtil.build(chatBO.getMessage(), chatBO.getPromptText()));
        }
        return agent.chatStream(chatBO.getConversationId(), chatBO.getPrompt().text());
    }


}
